import torch.nn as nn
import torch.nn.functional as F
class PositionwiseFeedForward(nn.Module):
    '''
        前馈神经网络
    '''

    def __init__(self, d_in, d_hid, dropout=0.1):

        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Conv1d(d_in, d_hid, 1)
        self.w_2 = nn.Conv1d(d_hid, d_in, 1)
        self.layer_normal = nn.LayerNorm(d_in)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        residual = x
        output = x.transpose(1, 2)
        output = self.w_2(F.relu(self.w_1(output)))
        output = output.transpose(1, 2)
        output = self.dropout(output)
        output = self.layer_normal(output + residual)
        return output


